import json
import os

from option import parse_args
from ffpp.preprocessing.data_split_ffpp import save_file


def get_videos(splits_path, mode):
    assert mode in ['train', 'val', 'test']
    videos = []
    json_path = os.path.join(splits_path, f'{mode}.json')
    with open(json_path, 'r') as dataset_json:
        dataset = json.load(dataset_json)
    for video_src, video_target in dataset:
        videos.append(f'{video_src}.mp4')
        videos.append(f'{video_target}.mp4')
    return videos


def main():
    args = parse_args()
    splits_path = '../../data/splits'
    assert os.path.exists(splits_path)
    train_video_names = get_videos(splits_path, mode='train')
    val_video_names = get_videos(splits_path, mode='val')
    test_video_names = get_videos(splits_path, mode='test')
    print(len(train_video_names), len(val_video_names), len(test_video_names))

    train_video_names.sort()
    val_video_names.sort()
    test_video_names.sort()
    save_file(train_video_names, splits_path, 'train.txt')
    save_file(val_video_names, splits_path, 'val.txt')
    save_file(test_video_names, splits_path, 'test.txt')


if __name__ == '__main__':
    # Total: 2000
    main()
